"""
面向 multipath-tools 软件的 MCP Server

MCP tools 列表：
1. kpartx_manage - 管理分区映射表操作
   用途：执行 kpartx 命令的各种操作，包括添加、更新、列出、删除分区映射
   参数：device (str): 设备路径, action (str): 操作类型, options (str): 额外选项
   返回：统一 JSON 结构

2. mpathconf_configure - 配置 multipath 工具
   用途：配置 multipath 的各种设置，包括启用/禁用、用户友好名称等
   参数：action (str): 配置操作, options (str): 配置选项
   返回：统一 JSON 结构

3. mpathpersist_manage - 管理持久化预留
   用途：执行 mpathpersist 命令的各种持久化预留操作
   参数：device (str): 设备路径, action (str): 操作类型, options (str): 预留参数
   返回：统一 JSON 结构

4. multipath_scan - 扫描 multipath 设备
   用途：扫描和显示 multipath 设备信息，包括详细列表、移除设备等
   参数：action (str): 扫描操作, device (str): 设备路径, options (str): 扫描选项
   返回：统一 JSON 结构

5. multipath_configure - 配置 multipath 参数
   用途：配置 multipath 的各种参数，包括重新配置、绑定设置等
   参数：action (str): 配置操作, device (str): 设备路径, options (str): 配置选项
   返回：统一 JSON 结构

6. multipathd_service - 管理 multipathd 服务
   用途：管理 multipathd 服务的启动、停止、重启和状态检查
   参数：action (str): 服务操作
   返回：统一 JSON 结构

7. multipathd_socket - 管理 multipathd socket
   用途：管理 multipathd socket 的执行和重载
   参数：action (str): socket操作
   返回：统一 JSON 结构
"""

from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json

mcp = FastMCP("Multipath-tools MCP Server")

@mcp.tool()
def kpartx_manage(device: str, action: str, options: str = "") -> dict:
    """
    执行 kpartx 命令管理分区映射表
    
    参数:
        device: 设备路径，如 /dev/sda
        action: 操作类型，如 'add', 'update', 'list', 'delete', 'show' 等
        options: 额外选项字符串，如 '-a -f -v' 等
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    try:
        if action == "add":
            cmd = f"kpartx -a {options} {device}"
        elif action == "update":
            cmd = f"kpartx -u {device}"
        elif action == "list":
            cmd = f"kpartx -l {device}"
        elif action == "delete":
            cmd = f"kpartx -d {device}"
        elif action == "show":
            cmd = f"kpartx -s {device}"
        else:
            cmd = f"kpartx {options} {device}"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"kpartx {action} {device}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def mpathconf_configure(action: str, options: str = "") -> dict:
    """
    配置 multipath 工具的各种设置
    
    参数:
        action: 配置操作，如 'enable', 'disable', 'set_user_friendly_names', 'set_find_multipaths' 等
        options: 配置选项字符串，如 '--allow 360000000000000000e00000000010001 --outfile multipath.conf' 等
    
    返回:
        统一 JSON 结构
    """
    try:
        if action == "enable":
            cmd = f"mpathconf --enable {options}"
        elif action == "disable":
            cmd = "mpathconf --disable"
        elif action == "set_user_friendly_names":
            cmd = "mpathconf --user_friendly_names y"
        elif action == "set_find_multipaths":
            cmd = "mpathconf --find_multipaths y"
        elif action == "status":
            cmd = "mpathconf"
        else:
            cmd = f"mpathconf {options}"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"mpathconf {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def mpathpersist_manage(device: str, action: str, options: str = "") -> dict:
    """
    管理持久化预留操作
    
    参数:
        device: 设备路径，如 /dev/mapper/mpath0
        action: 操作类型，如 'register', 'query', 'reserve', 'release' 等
        options: 预留参数字符串，如 '--param-sark=12cdbe' 等
    
    返回:
        统一 JSON 结构
    """
    try:
        if action == "register":
            cmd = f"mpathpersist --out --register {options} {device}"
        elif action == "query":
            cmd = f"mpathpersist -i -k {options} {device}"
        elif action == "reserve":
            cmd = f"mpathpersist --out --reserve {options} {device}"
        elif action == "release":
            cmd = f"mpathpersist --out --release {options} {device}"
        else:
            cmd = f"mpathpersist {options} {device}"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"mpathpersist {action} {device}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def multipath_scan(action: str, device: str = "", options: str = "") -> dict:
    """
    扫描和显示 multipath 设备信息
    
    参数:
        action: 扫描操作，如 'list', 'list_detailed', 'remove', 'add', 'check' 等
        device: 设备路径（可选）
        options: 扫描选项字符串，如 '-v3' 等
    
    返回:
        统一 JSON 结构
    """
    try:
        if action == "list":
            cmd = "multipath -l"
        elif action == "list_detailed":
            cmd = "multipath -ll"
        elif action == "remove":
            cmd = f"multipath -v3 -f {device}"
        elif action == "add":
            cmd = f"multipath -a {device}"
        elif action == "check":
            cmd = f"multipath -v3 -C {device}"
        elif action == "verbose_scan":
            cmd = "multipath -v3"
        else:
            cmd = f"multipath {options} {device}"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"multipath {action} {device}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def multipath_configure(action: str, device: str = "", options: str = "") -> dict:
    """
    配置 multipath 的各种参数
    
    参数:
        action: 配置操作，如 'reconfigure', 'show_config', 'set_policy', 'check_device' 等
        device: 设备路径（可选）
        options: 配置选项字符串，如 '-p multibus' 等
    
    返回:
        统一 JSON 结构
    """
    try:
        if action == "reconfigure":
            cmd = "multipath -r -v3"
        elif action == "show_config":
            cmd = "multipath -t"
        elif action == "set_policy":
            cmd = f"multipath -v3 -p {options}"
        elif action == "check_device":
            cmd = f"multipath -c {device}"
        elif action == "reset_wwids":
            cmd = f"multipath -W {device}"
        elif action == "remove_wwid":
            cmd = f"multipath -w {device}"
        else:
            cmd = f"multipath {options} {device}"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"multipath {action} {device}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def multipathd_service(action: str) -> dict:
    """
    管理 multipathd 服务
    
    参数:
        action: 服务操作，如 'start', 'stop', 'restart', 'status', 'reload' 等
    
    返回:
        统一 JSON 结构
    """
    try:
        if action == "start":
            cmd = "systemctl start multipathd.service"
        elif action == "stop":
            cmd = "systemctl stop multipathd.service"
        elif action == "restart":
            cmd = "systemctl restart multipathd.service"
        elif action == "status":
            cmd = "systemctl status multipathd.service"
        elif action == "reload":
            cmd = "systemctl reload multipathd.service"
        else:
            cmd = f"systemctl {action} multipathd.service"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"multipathd_service {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def multipathd_socket(action: str) -> dict:
    """
    管理 multipathd socket
    
    参数:
        action: socket操作，如 'start', 'stop', 'restart', 'reload' 等
    
    返回:
        统一 JSON 结构
    """
    try:
        cmd = f"systemctl {action} multipathd.socket"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"multipathd_socket {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()